# -*- coding:utf8 -*-

from matplotlib.pylab import *
from layers import *
from load_data import *

import theano
import Image
import datetime


class ParamDumpHelper(object):
    """
    会将每一次的epoch对应的参数保存，包括W,b，train_sore,valid_sore
    """
    def __init__(self,dump_file):
        self.dump_file = dump_file

    def get_params_from_file(self):
        if os.path.exists(self.dump_file):
            f = gzip.open(self.dump_file,"rb")
            dumped_file = pickle.load(f)
            f.close()
            return dumped_file
        return []

    def params_to_file(self,params):
        f = gzip.open(self.dump_file,"wb")
        if len(params) > 20:
            params = params[7:]
        pickle.dump(params,f)
        f.close()

    def dump(self, params):
        dumped_params = self.get_params_from_file()
        dumped_params.append(params)
        self.params_to_file(dumped_params)

class Spp_detect(object):
    def __init__(self, exist_params,batch_size):
        self.rng = numpy.random.RandomState(1234)
        self.exist_params = exist_params
        self.batch_size = batch_size

    def layer_params(self, nkerns,x_dim,y_dim):
        src_channel = 3
        self.layer1_image_shape  = (self.batch_size, src_channel, x_dim, y_dim)
<<<<<<< HEAD
        self.layer1_filter_shape = (nkerns[0], src_channel, 7, 7)
        self.layer2_image_shape  = (self.batch_size, nkerns[0], x_dim-6, y_dim-6)
        self.layer2_filter_shape = (nkerns[1], nkerns[0], 7, 7)
        self.layer3_image_shape  = (self.batch_size, nkerns[1], x_dim-12, y_dim-12)
        self.layer3_filter_shape = (nkerns[2], nkerns[1], 7, 7)
        self.layer4_image_shape  = (self.batch_size, nkerns[2], x_dim-18, y_dim-18)
        self.layer4_filter_shape = (nkerns[3], nkerns[2], 7, 7)
        self.layer5_image_shape  = (self.batch_size, nkerns[3], x_dim-24, y_dim-24)
        self.layer5_filter_shape = (nkerns[4], nkerns[3], 7, 7)
        self.result_image_shape  = (self.batch_size, nkerns[4], x_dim-30, y_dim-30)
=======
        self.layer1_filter_shape = (nkerns[0], src_channel, 3, 3)
        self.layer2_image_shape  = (self.batch_size, nkerns[0], x_dim-2, y_dim-2)
        self.layer2_filter_shape = (nkerns[1], nkerns[0], 3, 3)
        self.layer3_image_shape  = (self.batch_size, nkerns[1], x_dim-4, y_dim-4)
        self.layer3_filter_shape = (nkerns[2], nkerns[1], 3, 3)
        self.layer4_image_shape  = (self.batch_size, nkerns[2], x_dim-6, y_dim-6)
        self.layer4_filter_shape = (nkerns[3], nkerns[2], 3,3 )
        self.layer5_image_shape  = (self.batch_size, nkerns[3], x_dim-8, y_dim-8)
        self.layer5_filter_shape = (nkerns[4], nkerns[3], 3, 3)
        self.result_image_shape  = (self.batch_size, nkerns[4], x_dim-10, y_dim-10)
>>>>>>> 8568c94926e3d0d8c2cb51272715f91c2e7295f8

    def build_layer_architecture(self, acti_func=relu):
        '''
        simple means the deepid layer input is only the layer4 output.
        layer1: convpool layer
        layer2: convpool layer
        layer3: convpool layer
        layer4: conv layer
        deepid: hidden layer
        softmax: logistic layer
        '''

        x = T.matrix('x')

        print 'building the model ...'

        layer1_input = x.reshape(self.layer1_image_shape)
        self.layer1 = ConvLayer(self.rng,
                input        = layer1_input,
                image_shape  = self.layer1_image_shape,
                filter_shape = self.layer1_filter_shape,
                W = self.exist_params[6][0],
                b = self.exist_params[6][1],
                activation   = acti_func)

        self.layer2 = ConvLayer(self.rng,
                input        = self.layer1.output,
                image_shape  = self.layer2_image_shape,
                filter_shape = self.layer2_filter_shape,
                W = self.exist_params[5][0],
                b = self.exist_params[5][1],
                activation   = acti_func)

        self.layer3 = ConvLayer(self.rng,
                input        = self.layer2.output,
                image_shape  = self.layer3_image_shape,
                filter_shape = self.layer3_filter_shape,
                W = self.exist_params[4][0],
                b = self.exist_params[4][1],
                activation   = acti_func)

        self.layer4 = ConvLayer(self.rng,
                input        = self.layer3.output,
                image_shape  = self.layer4_image_shape,
                filter_shape = self.layer4_filter_shape,
                W = self.exist_params[3][0],
                b = self.exist_params[3][1],
                activation   = acti_func)

        self.layer5 = ConvLayer(self.rng,
                input        = self.layer4.output,
                image_shape  = self.layer5_image_shape,
                filter_shape = self.layer5_filter_shape,
                W = self.exist_params[2][0],
                b = self.exist_params[2][1],
                activation   = acti_func)

        spp_input = self.layer5.output

        self.spp_layer = SppLayer_train(
                input = spp_input,
                bins= [6,3,1]
                )
        self.feature_map = theano.function(inputs=[x],
                                           outputs=self.layer5.output)

    def get_feature_map(self,x):
        print "generating feature_map"
        starttime = datetime.datetime.now()
        feature_map = self.feature_map(x)
        endtime = datetime.datetime.now()
        print "time consuming: " + str((endtime - starttime).seconds)
        return feature_map
<<<<<<< HEAD
class predict(object):
    def __init__(self,exist_params):
        self.rng = numpy.random.RandomState(1234)
        self.exist_params = exist_params

    def build_layer_architecture(self, acti_func=relu,n_spp=0,n_hidden=0,n_out=0):
        print "building detect model ..."

        x = T.matrix('x')
        spp_input = x
        self.spp_layer = SppLayer_train(
                input = spp_input,
                bins= [4,2,1]
                )

        self.hidden_layer = HiddenLayer(self.rng,
                input = self.spp_layer.output,
                n_in  = n_spp,
                # n_in  = numpy.prod( self.result_image_shape[1:] ),
                n_out = n_hidden,
                W = self.exist_params[1][0],
                b = self.exist_params[1][1],
                activation = acti_func)

        self.softmax_layer = LogisticRegression(
                input = self.hidden_layer.output,
                n_in  = n_hidden,
                n_out = n_out,
                W = self.exist_params[0][0],
                b = self.exist_params[0][1])

        self.predict_result=theano.function(inputs=[x],outputs=self.softmax_layer.getPred())

    def predict_func(self, x):
        print '\tpredicting ...'
        result_judge,p_y_given_x = self.predict_result(x)
        return result_judge,p_y_given_x
=======
>>>>>>> 8568c94926e3d0d8c2cb51272715f91c2e7295f8

def window_to_feature_map(window_point_x1,window_point_y1,window_point_x2,window_point_y2,
                          window_size_x,window_size_y,map_size_x,map_size_y):
    map_point_x1 = window_point_x1*math.ceil(map_size_x/window_size_x)-1
    map_point_y1 = window_point_y1*math.ceil(map_size_y/window_size_y)-1
    map_point_x2 = window_point_x2*math.ceil(map_size_x/window_size_x)-1
    map_point_y2 = window_point_y2*math.ceil(map_size_y/window_size_y)-1

    return map_point_x1,map_point_y1,map_point_x2,map_point_y2

def image_to_featrue_map(params_file,image_path,x_dim,y_dim):
    arrs = []
    image_vector_len = y_dim * x_dim * 3
    im = Image.open(image_path)
    im = np.asarray(im, dtype='string')
    arr_img = im.reshape((image_vector_len,))
    arrs.append(arr_img)
    arrs = np.asarray(arrs, dtype='float32')
    arrs = np.vstack(arrs)

    spp_detect = build_model(params_file=params_file,x_dim=x_dim,y_dim=y_dim)
    feature_map = spp_detect.get_feature_map(arrs)

    return feature_map

def build_model(params_file,x_dim,y_dim):
    pd_helper = ParamDumpHelper(params_file)
    exist_params = pd_helper.get_params_from_file()
    if len(exist_params) != 0:
        exist_params = exist_params[-1]
    else:
        print 'error, no trained params'
        return

    spp_detect = Spp_detect(exist_params=exist_params,batch_size=1)
    spp_detect.layer_params(nkerns=[20,20,40,40,80],x_dim=x_dim,y_dim=y_dim)
    spp_detect.build_layer_architecture(acti_func=relu)
<<<<<<< HEAD
=======

>>>>>>> 8568c94926e3d0d8c2cb51272715f91c2e7295f8
    return spp_detect

def feature_map_to_file(feature_map,feature_map_path):
    f = gzip.open(feature_map_path, 'wb')
    pickle.dump(feature_map, f)
    f.close()


def show_feature_map(params_file,image_path,x_dim,y_dim):
    feature_map = image_to_featrue_map(params_file=params_file,image_path=image_path,x_dim=x_dim,y_dim=y_dim)
    flag = 1
    for i in range(4):
        subplot(2,2,flag)
        imshow(feature_map[0][i])
        flag += 1
        axis("off")
    show()

<<<<<<< HEAD
def detect(feature_map_path,param_path,n_spp,n_hidden,n_out,acti_func):
    f = gzip.open(feature_map_path, 'rb')
    feature_map = pickle.load(f)
    f.close()

    PFHelper = ParamDumpHelper(param_path)
    exist_params = PFHelper.get_params_from_file()
    if len(exist_params) != 0:
        print "no exist params"
        return

    net = predict(exist_params=exist_params)
    net.build_layer_architecture(acti_func=relu,n_spp=1680,n_hidden=160,n_out=2)





def crop_patch(src_file_path,res_file_path,patch_dim):
    image_vector_len=patch_dim*patch_dim*3
    im = Image.open(src_file_path)
    im_draw=ImageDraw.Draw(im)
    x_size,y_size=im.size
    x_one=0
    y_one=0
    x_two=patch_dim
    y_two=patch_dim
    flag=1
    detect_patch=teaDetect64_deeper.judge(params_file="C:\\Users\\wuxiaomin\\Desktop\\arobei\\new\\param_e40_001_cpu"\
                                  ,n_hidden=160,n_out=2,acti_func=relu)
    stride = 32
    for i in range(int((y_size-patch_dim)/stride)):
        arrs=[]
        for j in range(int((x_size-patch_dim)/stride)):
            box=(x_one+j*stride,y_one+i*stride,x_two+j*stride,y_two+i*stride)
            img_crop=im.crop(box)
            if ((x_two+j)-(x_one+j))!=patch_dim and ((y_two+i)-(y_one+i))!=patch_dim:
                print "crop wrong"
                return
            arr_img = np.asarray(img_crop, dtype='string')
            arr_img = arr_img.transpose(2,0,1).reshape((image_vector_len,))
            arrs.append(arr_img)
        print "num h:"+str(flag)
        arrs = np.asarray(arrs, dtype='float32')
        arrs = np.vstack(arrs)
        global batch_size_input
        batch_size_input = len(arrs)
        Predict,Propa=detect_patch.result_judgefunc(arrs[0:48])
        print Predict
        for index,value in enumerate(Predict):
            if Propa[index][1]>0.6:
                im_draw.line(((index*stride,flag*stride), (index*stride+patch_dim,flag*stride)),(255,0,0),5)
                im_draw.line(((index*stride,flag*stride), (index*stride,flag*stride+patch_dim)),(255,0,0),5)
                im_draw.line(((index*stride,flag*stride+patch_dim), (index*stride+patch_dim,flag*stride+patch_dim)),(255,0,0),5)
                im_draw.line(((index*stride+patch_dim,flag*stride), (index*stride+patch_dim,flag*stride+patch_dim)),(255,0,0),5)
        flag +=1
    im.save(res_file_path)



if __name__ == '__main__':
    show_feature_map(params_file="/home/eelab/wxm/data/tea/result/param_e20_l0001_h160_spp",
                     image_path="/home/eelab/wxm/data/tea/CY (458)-half.jpg",
=======
if __name__ == '__main__':
    show_feature_map(params_file="C:\\Users\\wuxiaomin\\Desktop\\arobei\\new\\param_e4_0001_cpu_spp",
                     image_path="C:\\Users\\wuxiaomin\\Desktop\\arobei\\tea_detect\\src\\CY (458)-half.jpg",
>>>>>>> 8568c94926e3d0d8c2cb51272715f91c2e7295f8
                     x_dim=1632,y_dim=1224)